import logging
import time
import zmq
from zmq.eventloop.ioloop import IOLoop, PeriodicCallback
from zmq.eventloop.zmqstream import ZMQStream

from kvsimple import KVMsg
from zhelpers import dump


class Route:
    def __init__(self, socket, identity, subtree):
        self.socket = socket
        self.identity = identity
        self.subtree = subtree


def send_single(key, kvmsg, route):
    if kvmsg.key.startswith(route.subtree):
        route.socket.send(route.identity, zmq.SNDMORE)
        kvmsg.send(route.socket)


class CloneServer(object):
    ctx = None
    kvmap = None
    loop = None
    port = None
    sequence = 0
    snapshot = None
    publisher = None
    collector = None

    def __init__(self, port=5556):
        self.port = port
        self.ctx = zmq.Context()
        self.kvmap = {}
        self.loop = IOLoop.instance()

        self.snapshot = self.ctx.socket(zmq.ROUTER)
        self.publisher = self.ctx.socket(zmq.PUB)
        self.collector = self.ctx.socket(zmq.PULL)
        self.snapshot.bind("tcp://*:%d" % self.port)
        self.publisher.bind("tcp://*:%d" % (self.port + 1))
        self.collector.bind("tcp://*:%d" % (self.port + 2))

        self.snapshot = ZMQStream(self.snapshot)
        self.publisher = ZMQStream(self.publisher)
        self.collector = ZMQStream(self.collector)

        self.snapshot.on_recv(self.handle_snapshot)
        self.collector.on_recv(self.handle_collect)
        self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)

        logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

    def start(self):
        self.flush_callback.start()
        try:
            self.loop.start()
        except KeyboardInterrupt:
            pass

    def handle_snapshot(self, msg):
        if len(msg) != 3 or msg[1] != b"ICANHAZ?":
            print("E: bad request, aborting")
            dump(msg)
            self.loop.stop()
        identity, request, subtree = msg
        if subtree:
            route = Route(self.snapshot, identity, subtree)

            for k, v in self.kvmap.items():
                send_single(k, v, route)

            logging.info("I: Sending state snapshot=%d" % self.sequence)
            self.snapshot.send(identity, zmq.SNDMORE)
            kvmsg = KVMsg(self.sequence)
            kvmsg.key = "KTHXBAI"
            kvmsg.body = subtree.decode('utf-8')
            kvmsg.send(self.snapshot)

    def handle_collect(self, msg):
        kvmsg = KVMsg.from_msg(msg)
        self.sequence += 1
        kvmsg.sequence = self.sequence
        kvmsg.send(self.publisher)
        ttl = float(kvmsg.get(b'ttl', 0))
        if ttl:
            kvmsg[b'ttl'] = b'%f' % (time.time() +ttl)
        kvmsg.store(self.kvmap)
        logging.info("I: publishing update=%d" % self.sequence)

    def flush_ttl(self):
        for key, kvmsg in list(self.kvmap.items()):
            self.flush_single(kvmsg)

    def flush_single(self, kvmsg):
        ttl = float(kvmsg.get(b'ttl', 0))
        if ttl and ttl <= time.time():
            kvmsg.body = ""
            self.sequence += 1
            kvmsg.sequence = self.sequence
            kvmsg.send(self.publisher)
            del self.kvmap[kvmsg.key]
            logging.info("I: publishing delete=%d", self.sequence)


def main():
    clone = CloneServer()
    clone.start()


if __name__ == "__main__":
    main()
